package client

import (
	"github.com/golang/glog"
	"time"
	"cloud_worm/util/proto/ujson"
	"cloud_worm/util/uid"
	"os"
	"fmt"
	"io/ioutil"
	"runtime"
	"strings"
	"cloud_worm/util/uexec"
	"cloud_worm/util/uconst"
	"cloud_worm/util/ufile"
	"errors"
)

func (a *Agent) PushMsg() {
	glog.Warningf("push msg Goroutine start...")
	if a.HeartBeatTimeout == 0 {
		a.HeartBeatTimeout = 60
	}
	t := time.NewTimer(time.Second * a.HeartBeatTimeout)

	var sendMsg = ujson.ProtoMsg{Type: ujson.ProtoHeartBeat}
	htPayload, _ := sendMsg.MarshalJSON()

	for {
		select {
		case <- t.C:
			a.TlsClient.Conn.Write(htPayload)
			t.Reset(time.Second * a.HeartBeatTimeout)

		case payload := <- a.PushMsgC:
			a.TlsClient.Conn.Write(payload)
			t.Reset(time.Second * a.HeartBeatTimeout)

		case <- a.CloseC:
			glog.Warningf("push msg Goroutine stop...")
			return
		}
	}

	return
}

func (a *Agent) ReqAgentId() error {
	var idMsg ujson.MsgId
	if len(a.Id) == 0 {
		hostId := uid.GetAgentIdFromFile(a.IdFilePath)
		if len(hostId) == 0 {
			// 读取失败,设置haveid为false，告诉服务器，agent是新的，没有id
			idMsg.HaveId = false
			idMsg.Id = ""
			glog.Infof("read file failed: %s", a.IdFilePath)
		} else {
			idMsg.HaveId = true
			idMsg.Id = string(hostId)
		}
	} else {
		idMsg.HaveId = true
		idMsg.Id = a.Id
	}
	idMsg.Pid = os.Getpid()

	// 协议数据包的封装
	var err error
	if content, err1 := idMsg.MarshalJSON(); err1 == nil {
		var sendMsg = ujson.ProtoMsg{
			Type: ujson.ProtoId,
			Content: content,
		}
		if buf, err2 := sendMsg.MarshalJSON(); err2 == nil {
			if err3 := a.TlsClient.Conn.Write(buf); err3 != nil {
				err = fmt.Errorf("send id msg error: %s", err3.Error())
			} else {
				err = nil
			}
		} else {
			err = fmt.Errorf("marshal id send msg error: %s", err2.Error())
		}
	} else {
		err = fmt.Errorf("marshal id msg error: %s", err1.Error())
	}

	return err
}

var (
	DropByServerErr = errors.New("dropped by server")
)

func (a *Agent) HandleAgentIdMsg() error {
	data, err := a.TlsClient.Conn.Read()
	if err != nil {
		a.TlsClient.Close()
		return DropByServerErr
	}

	var protoMsg ujson.ProtoMsg
	if err := protoMsg.UnmarshalJSON(data); err != nil {
		return err
	}
	if protoMsg.Type != ujson.ProtoId {
		return fmt.Errorf("proto id msg type error")
	}

	var idMsg ujson.MsgId
	if err := idMsg.UnmarshalJSON(protoMsg.Content); err != nil {
		return err
	}

	if idMsg.HaveId == false {
		return fmt.Errorf("server check id error")
	}
	if len(idMsg.Id) == 0 {
		return fmt.Errorf("server create id failed")
	}

	// 记录到实例
	a.Id = idMsg.Id

	/* 首次创建，写入/etc/worm/xxx文件中 */
	if idMsg.IsCreate == true {
		dir, err := ioutil.ReadDir(a.IdFilePath)
		if err != nil {
			glog.Errorf("open worm conf file error!")
		}
		// 该目录的全部文件删除文件
		for _, item := range dir {
			os.Remove(a.IdFilePath + "/" + item.Name())
		}

		var tmp = []byte(idMsg.Id)
		err = ioutil.WriteFile(a.IdFilePath + "/" + a.Id, tmp, 0600)
		if err != nil {
			return fmt.Errorf("local write File error: %s", err.Error())
		}
	}

	return nil
}

func (a *Agent) SendVersion() {
	versionMsg := ujson.MsgVersion{
		Version: a.Version,
		Platform: runtime.GOOS,
	}
	content, err := versionMsg.MarshalJSON()
	if err != nil {
		glog.Errorf("version msg Marshal error: %s", err.Error())
		return
	}

	sendMsg := ujson.ProtoMsg{
		Type: ujson.ProtoVersion,
		Content: content,
	}
	buf, err := sendMsg.MarshalJSON()
	if err != nil {
		glog.Errorf("send msg Marshal error: %s", err.Error())
		return
	}

	err = a.TlsClient.Conn.Write(buf)
	if err != nil {
		glog.Errorf("send version msg error: %s", err.Error())
		return
	}

	return
}

func (a *Agent) HandleMsg(localIP string) {
	for {
		data, err := a.TlsClient.Conn.Read()
		if err != nil {
			a.TlsClient.Close()

			a.CloseC <- struct{}{}
			glog.Errorf("conn dropped: %s", err.Error())

			a.TryLogin()
			continue
		}

		var protoMsg ujson.ProtoMsg
		if err := protoMsg.UnmarshalJSON(data); err != nil {
			glog.Errorf("proto msg Unmarshal error: %s", err.Error())
			continue
		}

		switch protoMsg.Type {
		case ujson.ProtoId:
			// id响应包，该逻辑在TryLogin()内部实现处理
		case ujson.ProtoCmdMsg:
			a.HandleCmdMsg(protoMsg.Content)
		case ujson.ProtoWpMsgReq:
			a.HandleWpMsg(protoMsg.Content)
		case ujson.ProtoAutoConfigReq:
			a.HandleAutoConfig(protoMsg.Content, localIP)
		case ujson.ProtoHeartBeat:
		default:
			glog.Warningf("unknown proto type: %d", protoMsg.Type)
		}
	}
}

func (a *Agent) HandleCmdMsg(data []byte)  {
	var cmdMsg ujson.MsgCmd
	if err := cmdMsg.UnmarshalJSON(data); err != nil {
		glog.Errorf("msg cmd Unmarshal error: %s", err.Error())
		return
	}

	// 处理命令
	if runtime.GOOS == "windows" {
		var inputParams []string
		if len(cmdMsg.Params) > 0 {
			inputParams = strings.Split(cmdMsg.Params, " ")
		}
		cmdMsg.Result, cmdMsg.Status = uexec.ExecWinCmd(cmdMsg.Cmd, inputParams, cmdMsg.Timeout)
		glog.Infof("Cmd: %s %s", cmdMsg.Cmd, cmdMsg.Params)
	} else if runtime.GOOS == "linux" || runtime.GOOS == "darwin" {
		cmd := cmdMsg.Cmd + " " + cmdMsg.Params
		cmdMsg.Result, cmdMsg.Status = uexec.ExecUnixCmd(cmd, cmdMsg.Timeout)
		glog.Infof("Cmd: %s", cmd)
	}
	glog.Infof("Output: %s", cmdMsg.Result)

	// 回复
	if content, err1 := cmdMsg.MarshalJSON(); err1 == nil {
		var sendMsg = ujson.ProtoMsg{
			Type: ujson.ProtoCmdMsg,
			Content: content,
		}
		if buf, err2 := sendMsg.MarshalJSON(); err2 == nil {
			if err3 := a.TlsClient.Conn.Write(buf); err3 != nil {
				glog.Warningf("dropped: %s", err3.Error())
				a.TryLogin()
			}
		} else {
			glog.Errorf("Cmd: %s, send msg Marshal error: %s", cmdMsg.Cmd, err1.Error())
			return
		}
	} else {
		glog.Errorf("Cmd: %s, msg Marshal error: %s", cmdMsg.Cmd, err1.Error())
		return
	}

	return
}

func (a *Agent) HandleWpMsg(data []byte) {
	var wpMsg ujson.MsgWpReq
	if err := wpMsg.UnmarshalJSON(data); err != nil {
		glog.Errorf("wp req msg Unmarshal error: %s", err.Error())
		return
	}

	if len(wpMsg.Id) != uconst.AgentIdLen || len(wpMsg.TaskId) <= 0 || wpMsg.WpId <= 0 {
		glog.Error("MsgWpReq params values error.")
		return
	}

	var respMsg = ujson.MsgWpResp{
		Id: wpMsg.Id,
		TaskId: wpMsg.TaskId,
		WpId: wpMsg.WpId,
		Input: wpMsg.Params,
	}

	/* 拼装目标文件 */
	var destFile string
	if runtime.GOOS == "windows" {
		destFile = WinBaseWorkDir + "\\" + WeaponDirName + "\\" + wpMsg.WpName + ".exe"
	} else if (runtime.GOOS == "linux") || (runtime.GOOS == "darwin") {
		destFile = UnixBaseWorkDir + "/" + WeaponDirName + "/" + wpMsg.WpName
	}

	md5Match := false
	if ufile.FileExists(destFile) == false {
		// local不存在destination file，则download file
		if err := ufile.DownloadFile(wpMsg.WpUrl, destFile); err != nil {
			glog.Errorf("url: '%s', destfile: '%s', failed: %s", wpMsg.WpUrl, destFile, err.Error())

			respMsg.Status = 1
			respMsg.Output = "get file failed"
			a.BackWpRespMsg(&respMsg)
			return
		}
	} else {
		// local存在destination file, 则先match md5
		if ufile.GetFileMd5(destFile) != wpMsg.WpMd5 {
			// md5不匹配，则删除，重新download file
			os.Remove(destFile)

			err := ufile.DownloadFile(wpMsg.WpUrl, destFile)
			if err != nil {
				glog.Errorf("url: '%s', destfile: '%s', failed: %s", wpMsg.WpUrl, destFile, err.Error())

				respMsg.Status = 1
				respMsg.Output = "Get file failed"
				a.BackWpRespMsg(&respMsg)
				return
			}
		} else {
			md5Match = true
		}
	}

	// 匹配md5值
	if md5Match == false {
		if ufile.GetFileMd5(destFile) != wpMsg.WpMd5 {
			os.Remove(destFile)

			glog.Errorf("the md5 value is not match: %s", destFile)
			respMsg.Status = 1
			respMsg.Output = "The Md5 value is Not match."
			a.BackWpRespMsg(&respMsg)
			return
		}
	}

	// chmod
	err := ufile.ChangeFileMode(destFile, 0700)
	if err != nil {
		os.Remove(destFile)
		glog.Errorf("'%s' Change mode failed: %s", destFile, err.Error())

		respMsg.Status = 1
		respMsg.Output = "Change mode failed."
		a.BackWpRespMsg(&respMsg)
		return
	}

	// deal params && exec
	if runtime.GOOS == "windows" {
		var inputParams []string
		if len(wpMsg.Params) > 0 {
			inputParams = strings.Split(wpMsg.Params, " ")
		}
		respMsg.Output, respMsg.Status = uexec.ExecWinCmd(destFile, inputParams, wpMsg.Timeout)
		glog.Info("Weapon: ", destFile + " " + wpMsg.Params)

	} else if runtime.GOOS == "linux" || runtime.GOOS == "darwin" {
		cmd := destFile + " " + wpMsg.Params
		respMsg.Output, respMsg.Status = uexec.ExecUnixCmd(cmd, wpMsg.Timeout)
		glog.Info("Weapon: ", cmd)
	}
	glog.Info("Output: ", respMsg.Output)

	a.BackWpRespMsg(&respMsg)

	return
}

func (a *Agent) BackWpRespMsg(respMsg *ujson.MsgWpResp) {
	content, err := respMsg.MarshalJSON()
	if err != nil {
		glog.Error("WeaponRespMsg Marshal failed.")
		return
	}

	var sendMsg = ujson.ProtoMsg{
		Type: ujson.ProtoWpMsgResp,
		Content: content,
	}
	buf, err := sendMsg.MarshalJSON()
	if err != nil {
		glog.Error("ProtoMsg(weapon resp) Marshal failed.")
		return
	}

	err = a.TlsClient.Conn.Write(buf)
	if err != nil {
		glog.Warning("Send weapon resp msg failed. dropped: ", err.Error())
		a.TryLogin()
	}
	return
}